#
# _analysisFunctions.R
#
# Functions for use in analysis.
#

makeStats <- function(regs,by.var="treatment.assign") { 
  # Function to make a default set of turnout stats.
  out <- regs[,list(voted=sum(yvar),count=sum(!is.na(yvar))),
               by=by.var]
  # Calculate turnout rate and SE.
  out[,voted.rate := voted/count]
  out[,voted.se   := sqrt( (voted.rate * (1-voted.rate))/count )]
  out
}

.simpleCap <- function(x) {
    s <- strsplit(x, " ")[[1]]
    paste(toupper(substring(s, 1, 1)), substring(s, 2),
          sep = "", collapse = " ")
}

makeRow <- function(tab,condition) {
  # Make a latex row for this condition given table tab.  
  out <- paste(.simpleCap(condition),
              sprintf("&  %1.1f  &  %1.1f  &  %1.1f  &  %1.1f \\\\", # first row=turnout rate.
              100*tab[treatment.assign == condition & party == "All",voted.rate],
              100*tab[treatment.assign == condition & party == "NPP",voted.rate],
              100*tab[treatment.assign == condition & party == "DEM",voted.rate],
              100*tab[treatment.assign == condition & party == "REP",voted.rate]))
  # Second row = standard error.
  out <- c(out,sprintf("         & (%1.1f) & (%1.1f) & (%1.1f) & (%1.1f) \\\\",
              100*tab[treatment.assign == condition & party == "All",voted.se],
              100*tab[treatment.assign == condition & party == "NPP",voted.se],
              100*tab[treatment.assign == condition & party == "DEM",voted.se],
              100*tab[treatment.assign == condition & party == "REP",voted.se]))
  # Third row = sample size.
  out <- c(out,sprintf("         & %s & %s & %s & %s \\\\",
              prettyNum(tab[treatment.assign == condition & party == "All",count],","),
              prettyNum(tab[treatment.assign == condition & party == "NPP",count],","),
              prettyNum(tab[treatment.assign == condition & party == "DEM",count],","),
              prettyNum(tab[treatment.assign == condition & party == "REP",count],",")))
  out <- c(out,"\\hline")
  return(out)
}

makeRow2 <- function(tab,condition) {
  # Make a latex row for this condition given table tab. v2 includes OTH party after REP.
  out <- paste(.simpleCap(condition),
              sprintf("&  %1.1f  &  %1.1f  &  %1.1f  &  %1.1f  &  %1.1f \\\\", # first row=turnout rate.
              100*tab[treatment.assign == condition & party == "All",voted.rate],
              100*tab[treatment.assign == condition & party == "NPP",voted.rate],
              100*tab[treatment.assign == condition & party == "DEM",voted.rate],
              100*tab[treatment.assign == condition & party == "REP",voted.rate],
              100*tab[treatment.assign == condition & party == "OTH",voted.rate]))
  # Second row = standard error.
  out <- c(out,sprintf("         & (%1.1f) & (%1.1f) & (%1.1f) & (%1.1f) & (%1.1f) \\\\",
              100*tab[treatment.assign == condition & party == "All",voted.se],
              100*tab[treatment.assign == condition & party == "NPP",voted.se],
              100*tab[treatment.assign == condition & party == "DEM",voted.se],
              100*tab[treatment.assign == condition & party == "REP",voted.se],
              100*tab[treatment.assign == condition & party == "OTH",voted.se]))
  # Third row = sample size.
  out <- c(out,sprintf("         & %s & %s & %s & %s & %s \\\\",
              prettyNum(tab[treatment.assign == condition & party == "All",count],","),
              prettyNum(tab[treatment.assign == condition & party == "NPP",count],","),
              prettyNum(tab[treatment.assign == condition & party == "DEM",count],","),
              prettyNum(tab[treatment.assign == condition & party == "REP",count],","),
              prettyNum(tab[treatment.assign == condition & party == "OTH",count],",")))
  out <- c(out,"\\hline")
  return(out)
}

makeWithin <- function(regs) {
  regs.within <- copy(regs[,grep("yvar|treat|vote\\.14|block\\.num|any\\.letter",names(regs),value=T),with=F])
  regs.within[,block.vote.14 := mean(yvar,na.rm=T),by="block.num"]
  regs.within[,yvar := yvar - block.vote.14]
  regs.within[,block.treat.elec.info := mean(treat.elec.info,na.rm=T),by="block.num"]
  regs.within[,treat.elec.info := treat.elec.info - block.treat.elec.info]
  regs.within[,block.treat.toptwo.info := mean(treat.toptwo.info,na.rm=T),by="block.num"]
  regs.within[,treat.toptwo.info := treat.toptwo.info - block.treat.toptwo.info]
  regs.within[,block.treat.partisan := mean(treat.partisan,na.rm=T),by="block.num"]
  regs.within[,treat.partisan := treat.partisan - block.treat.partisan]
  regs.within[,block.any.letter := mean(any.letter,na.rm=T),by="block.num"]
  regs.within[,any.letter := any.letter - block.any.letter]
  regs.within
}

makeHetEffs <- function(regs,group.name="All") {
  # Create a data.table of treatment effects for a subgroup in pursuit of Figure 1.
  # V2 makes comparisons only to control group.
  within <- makeWithin(regs) # For within FE estimator.
  # First effect: any letter versus control.
  any.v.control <- within[,lm(I(100*yvar) ~ any.letter)]
  # Second and third effects: identity and info compared to placebo.
  trts.v.control <- within[,lm(I(100*yvar) ~ treat.toptwo.info + treat.partisan + treat.elec.info)]
  out <- data.table(coefs=c(coef(any.v.control)['any.letter'],
                            coef(trts.v.control)['treat.elec.info'],
                            coef(trts.v.control)['treat.toptwo.info'],
                            coef(trts.v.control)['treat.partisan']),
                    serr=sqrt(c(vcov(any.v.control)['any.letter','any.letter'],
                                vcov(trts.v.control)['treat.elec.info','treat.elec.info'],
                                vcov(trts.v.control)['treat.toptwo.info','treat.toptwo.info'],
                                vcov(trts.v.control)['treat.partisan','treat.partisan'])),
                    lab=factor(c("Any letter v control","Election info v control","Top-two info v control","Partisan v control"),levels=c("Any letter v control","Election info v control","Top-two info v control","Partisan v control")),
                    group=rep(group.name,4),pch.num=c(18,19,17,15))
  out
}

plotTreatEffs <- function(het.effs,leg.text=c("Any letter v control","Election info v control","Top-two info v control","Partisan v control"),leg.pch=c(19,17,15)) {
  # Plot treatment effect point estimates and confidence intervals.
  het.effs[,x.loc := seq_len(.N) + (as.numeric(group.fac) - 1)] # Location on x-axis, with offset by group.
  # Plot points.
  het.effs[,plot(x=x.loc,y=coefs,pch=pch.num,ylim=range(c(ub,lb,0)),ann=F,axes=F,cex=2)]
  abline(h=0,lty=2,lwd=2,col='gray') # zero line
  # Plot confidence intervals.
  het.effs[,segments(x0=x.loc,x1=x.loc,y0=lb,y1=ub,lwd=1)]
  axis(2,las=2)
  # Axis indicators for groups.
  grp.locs <- het.effs[,list(x.loc=mean(x.loc),x.min=min(x.loc),x.max=max(x.loc)),by=group]
  axis(1,at=grp.locs[,x.loc],labels=grp.locs[,group],tick=F,line=1)
  grp.locs[,axis(1,at=c(x.min,x.max),labels=FALSE),by=group]
  title(ylab="Conditional average treatment effect")
  # Legend.
  legend("topleft",legend=leg.text,pch=leg.pch)
}

makeCoefNames <- function(lm.list,coef.labs=list("(Intercept)"="Intercept",
                                  "any.letter"="Any letter","treat.elec.info"="Election info letter",
                                  "treat.toptwo.info"="Top-two info letter",
                                  "treat.partisan"="Partisan letter",
                                  "age.mean.dev"="Age","no.vote.08.10"="Abstain 08 and 10",
                                  "party.npp"="Party NPP","party.rep"="Party REP",
                                  "party.oth"="Party OTH",
                                  "factor(age.bin)2"="Age 29-38","factor(age.bin)3"="Age 39-48",
                                  "factor(age.bin)4"="Age 49-58","factor(age.bin)5"="Age 59-68",
                                  "factor(age.bin)6"="Age 69+",
                                  "reg.date.pre.08"="Registered prior to 08",
                                  "reg.date.pre.10"="Registered prior to 10")) {
  # Turn a list of lms into a coefficient name vector. Assumes coefs are listed
  # consistent with argument "lr" in apsrtable.
  coefs <- NULL
  for (i in 1:length(lm.list)) {
    coefs <- union(coefs,names(coef(lm.list[[i]])))
  }
  # Original names ordered appropriately in vector coefs. Now, new names.
  out <- NULL
  for (i in 1:length(coefs)) {
    if (coefs[i] %in% names(coef.labs)) { # exact name
      out <- c(out,coef.labs[[coefs[i]]])
    } else if (regexpr(":",coefs[i]) != -1) { # interaction term.
      splitter <- unlist(strsplit(coefs[i],":"))
      new.coef <- NULL
      for (j in 1:length(splitter)) { # see if each element has a match.
        if (splitter[j] %in% names(coef.labs)) {
          new.coef <- c(new.coef,coef.labs[[splitter[j]]])
        } else {
          new.coef <- c(new.coef,splitter[j])
        }
      }
      # Collapse with asterix.
      out <- c(out,paste(new.coef,collapse="*"))
    } else { # can't find match.
      out <- c(out,coefs[i])
    }
  }
  #print(cbind(coefs,out))
  return(out)
}
